package com.sicau.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.transformer.KeywordMetadataEnricher;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 文档加强处理器
 */
@Slf4j
@Component
public class KWMetadataEnricher {

    private final KeywordMetadataEnricher enricher;

    public KWMetadataEnricher(ChatModel chatModel) {
        enricher = KeywordMetadataEnricher.builder(chatModel)
                .keywordCount(5)
                .build();
    }

    /**
     * 添加关键词元数据
     * @param documents
     * @return
     */
    public List<Document> enrich(List<Document> documents) {
        log.info("文档加强");
        return enricher.apply(documents);
    }
}
